﻿#pragma once

#include "pch.h"
#include <string>
#include <unordered_map>
#include "RenderDefinitions.h"
#include "Common/NonCopyable.h"
#include "Utils/D3DUtils.h"

class FGlobalShaderMap;
class FShader;
class FShaderRegistration;
template<typename ShaderType>
class TShaderMapRef;

extern FGlobalShaderMap* ShaderMap;

class FGlobalShaderMap : public NonCopyable
{
public:
    FGlobalShaderMap() = default;
    ~FGlobalShaderMap() = default;

    template<typename ShaderType>
    ShaderType* GetShader(int32_t PermutationId) const;

    FShaderRegistration& FindOrAddRegistration(size_t InTypeHash);
    void AddShaderContent(uint64_t InShaderHash, Microsoft::WRL::ComPtr<ID3DBlob> InShaderBlob);

private:
    std::unordered_map<size_t, FShaderRegistration> MapRegistrations;
    std::unordered_map<uint64_t, Microsoft::WRL::ComPtr<ID3DBlob>> MapContent;
};

template<typename ShaderType>
class TShaderMapRef
{
public:
    TShaderMapRef(const FGlobalShaderMap* InShaderMap)
        : ShaderMap(InShaderMap)
    {
        check(ShaderMap);
        ShaderContent = ShaderMap->GetShader<ShaderType>(0);
    }

    ~TShaderMapRef()
    {
        delete ShaderContent;
        ShaderContent = nullptr;
    }

    BYTE* GetBufferPointer() const
    {
        return reinterpret_cast<BYTE*>(ShaderContent->GetBufferPointer());
    }
    
    SIZE_T GetBufferSize() const
    {
        return ShaderContent->GetBufferSize();
    }

private:
    const FGlobalShaderMap* ShaderMap;
    ShaderType* ShaderContent;
};

class FShader : public NonCopyable
{
public:
    FShader() = default;
    FShader(const Microsoft::WRL::ComPtr<ID3DBlob>& InShaderBlob)
        : ShaderBlob(InShaderBlob)
    {}

    BYTE* GetBufferPointer() const
    {
        return reinterpret_cast<BYTE*>(ShaderBlob->GetBufferPointer());
    }
    
    SIZE_T GetBufferSize() const
    {
        return ShaderBlob->GetBufferSize();
    }
    
private:
    Microsoft::WRL::ComPtr<ID3DBlob> ShaderBlob;
};

class FGlobalShader: public FShader
{
public:
    FGlobalShader() = default;
    FGlobalShader(const Microsoft::WRL::ComPtr<ID3DBlob>& InShaderBlob)
        : FShader(InShaderBlob)
    {}
};

class FShaderRegistration
{
public:
    void GenerateShaderFileNameHashId();
    uint64_t GetShaderFileNameHashId(int32_t PermutationId) const;

    inline void SetShaderName(const std::string& InShaderName) { ShaderName = InShaderName; }
    inline const std::string& GetShaderName() const { return ShaderName; }

    inline void SetShaderFilePath(const std::string& InShaderFilePath) { ShaderFilePath = InShaderFilePath; }
    inline const std::string& GetShaderFilePath() const { return ShaderFilePath; }

    inline void SetEntry(const std::string& InEntry) { Entry = InEntry; }
    inline const std::string& GetEntry() const { return Entry; }

    inline void SetFrequency(EShaderFrequency InFrequency) { Frequency = InFrequency; }
    inline EShaderFrequency GetFrequency() const { return Frequency; }
    
private:
    uint64_t ShaderFileHashId = 0;
    std::string ShaderName;
    std::string ShaderFilePath;
    std::string Entry;
    EShaderFrequency Frequency = SF_NumBits;
};

template<typename ShaderType>
class TGlobalShaderRegister
{
public:
    TGlobalShaderRegister(const std::string& InShaderName,
        const std::string& InShaderFilePath,
        const std::string& InEntry,
        const EShaderFrequency& InFrequency)
    {
        if (!ShaderMap)
        {
            ShaderMap = new FGlobalShaderMap();
        }
        
        size_t TypeHash = typeid(ShaderType).hash_code();

        FShaderRegistration& Registration = ShaderMap->FindOrAddRegistration(TypeHash);
        Registration.SetShaderName(InShaderName);
        Registration.SetShaderFilePath(InShaderFilePath);
        Registration.SetEntry(InEntry);
        Registration.SetFrequency(InFrequency);
        Registration.GenerateShaderFileNameHashId();

        uint64_t ShaderHash = Registration.GetShaderFileNameHashId(0);

        Microsoft::WRL::ComPtr<ID3DBlob> ShaderBlob = D3DUtil::CompileShader(std::wstring(InShaderFilePath.begin(), InShaderFilePath.end()),
            nullptr, InEntry, ConvertShaderFrequency(InFrequency, false));
        
        ShaderMap->AddShaderContent(ShaderHash, ShaderBlob);
    }
};

template<typename ShaderType>
ShaderType* FGlobalShaderMap::GetShader(int32_t PermutationId) const
{
    size_t TypeHash = typeid(ShaderType).hash_code();
    const FShaderRegistration& Registration = MapRegistrations.at(TypeHash);

    uint64_t ShaderHash = Registration.GetShaderFileNameHashId(PermutationId);
    return new ShaderType(MapContent.at(ShaderHash));
}

#define IMPLEMENT_GLOBAL_SHADER(Type, File, Entry, Frequency) static TGlobalShaderRegister<Type> Kite_Shader_Obj_##Type(#Type, File, Entry, Frequency);
